#!/usr/bin/env python3
# -*- coding: utf_8 -*-

import os
import sys
import time
import subprocess
import cmd
import traceback
import shlex
from pygments import highlight, lexers, formatters

from mie.bprint import cp, varfmt
from mie.roar.roar import Trans_Socket, CallPkg, CallDic, CallEvt


class MieCommandClient(cmd.Cmd):
    try:
        intro = subprocess.check_output(["figlet", "-f", "slant", "HEY, I IS M.CON"])
    except:
        intro = b"HEY, I IS MCON"
    intro = intro.decode("utf-8")

    prompt = cp.c("MCon$ ") + "\r\n"
    use_rawinput = True
    cmd.Cmd.identchars += "."

    def got_allcmds(self, ev, cookie, callctx, respdic):
        dic = respdic.content
        self.remote_name = dic.get("name", "MCon")
        self.setcmds(dic.get("cmds"))

    def on_conn_close(self, conn):
        self.clrcmds()

        self.trans = None
        self.prompt = cp.c("MCon$ ") + "\r\n"

    def connect(self, host, port):
        try:
            trans = Trans_Socket(None, host, port)

            callpkg = CallPkg.new(["hello"], ok=self.got_allcmds)
            callpkg.sendvia(trans)

            trans = trans
            conn = trans.conn

            conn.set_on_disconnect(self.on_conn_close)

            self.host = host
            self.port = port

            self.prompt = (
                cp.c("%s@%s:%d$ " % (cp.r(self.remote_name), host, port)) + "\r\n"
            )
        except:
            conn = None
            trans = None

        self.conn = conn
        self.trans = trans

    def complete_man(self, text, line, begidx, endidx):
        word = line[begidx:endidx]
        return [x for x in self.commands if x.startswith(word)]

    def __init__(self, host, port):
        cmd.Cmd.__init__(self)

        self.host = host
        self.port = port

        self.commands = set()
        self.connect(host, port)

    def precmd(self, line):
        self.args = shlex.split(line)
        return line

    def clrcmds(self):
        for c in self.commands:
            try:
                delattr(self.__class__, "do_" + c)
            except:
                pass

    def setcmds(self, cmds):
        for c in cmds or []:
            self.commands.add(c)

            try:
                setattr(self.__class__, "do_" + c, self.call_command)
            except:
                traceback.print_exc()

    def call_command(self, line):
        """Forward the command to target"""

        self.calldone = False

        def done(ev, cookie, *args):
            self.calldone = True

        def ok(ev, cookie, cmdctx, respdic):
            out = varfmt(respdic.content)
            self.last_output = out

            colorful_json = highlight(
                out.encode("UTF-8"), lexers.JsonLexer(), formatters.TerminalFormatter()
            )
            print(colorful_json)

        def ng(ev, cookie, cmdctx, respdic):
            print(vars(cmdctx))
            print("%s %s" % (cp.r("NG: Cmd:"), cmdctx.callpkg.calldic.name))
            print("%s %s" % (cp.r("NG: Err:"), respdic.errcode))
            print("%s %s" % (cp.r("NG: Msg:"), respdic.errmsg))

        if self.args[0] not in self.commands:
            return

        calldic = CallDic.fr_cmd(self.args)
        callevt = CallEvt()
        callevt.onok(ok).onng(ng).ondone(done).oncs(done).onfail(done)

        callpkg = CallPkg(calldic, callevt)
        callpkg.postvia(self.trans)

        loops = 0
        while not self.calldone:
            loops += 1
            time.sleep(0.01)
            if loops > 500:
                break

    def emptyline(self):
        if not self.trans:
            self.connect(self.host, self.port)

    def postloop(self):
        if self.trans:
            self.trans.bye()
        self.trans = None

    def do_s(self, line):
        """> s filename - Save last output to file"""

        fn = self.args[1] if len(self.args) > 1 else "/tmp/mcon.txt"
        try:
            with open(fn, "wt") as f:
                f.write(self.last_output)
            print("%s saved, %d bytes written" % (fn, len(self.last_output)))
        except:
            pass

    def do_conn(self, line):
        """> conn a [b] - (Re)connect to remote"""

        host, port = None, 0
        for arg in self.args[1:]:
            if port == 0 and arg.isdigit():
                port = int(arg)
            elif host is None:
                host = arg

        host = host or self.host
        port = port or self.port

        #
        # disconnect if necessary
        #
        if self.trans:
            if host == self.trans.conn.host and port == self.trans.conn.port:
                return

            print(
                "Will break the old connection to %s:%d"
                % (self.trans.conn.host, self.trans.conn.port)
            )
            self.trans.bye()
            self.trans = None

        self.connect(host, port)
        if self.trans:
            print(
                "Connect to %s:%d successfully."
                % (self.trans.conn.host, self.trans.conn.port)
            )
        else:
            print("Connect to %s:%d failed." % (host, port))

    def do_EOF(self, line):
        print("Bye")
        os._exit(0)


if __name__ == "__main__":
    argv = sys.argv[1:] + ["localhost", "9010"]

    host, port = None, None
    for arg in argv:
        if port is None and arg.isdigit():
            port = int(arg)
        elif host is None:
            host = arg

    MieCommandClient(host, port).cmdloop()
